from django.shortcuts import HttpResponse
import json
from langchain_community.document_loaders import TextLoader, PyPDFLoader, CSVLoader, Docx2txtLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.document_loaders import TextLoader
import os
from model.my_chat_model import ChatModel
from langchain.chains import RetrievalQA
from langchain_community.vectorstores import Neo4jVector
from langchain_neo4j import Neo4jGraph
import warnings
from langchain_core._api import LangChainDeprecationWarning

warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)


# 文档入库
def upload_file(request):
    # 文件上传
    if request.method == "POST":
        # 获取文件名
        file = request.FILES.get("file")
        print(f"file={file}")
        # 文件写入
        with open(f"./static/file/{file}", "wb") as f:
            for chunk in file.chunks():
                f.write(chunk)
        # 获取文件的后缀名
        file_type = file.name.split(".")[-1]
        print(f"file_type={file_type}")
        # 文档入库
        doc_store(file_type, file.name)

    return HttpResponse(json.dumps({"code": 200, "msg": "success"}))


# 文档存储
def doc_store(file_type, file_name):
    # 获取文件路径
    path = os.path.join("static", "file")
    print(f"path={path}")
    if file_type == "txt":
        loader = TextLoader(f"{path}/{file_name}", encoding="utf-8")
    elif file_type == "pdf":
        loader = PyPDFLoader(f"{path}/{file_name}")
    elif file_type == "csv":
        loader = CSVLoader(f"{path}/{file_name}", encoding="utf-8")
    else:
        loader = Docx2txtLoader(f"{path}/{file_name}")

    docs = loader.load_and_split()
    # 2 文本分割
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=100,
        chunk_overlap=50
    )
    split_docs = text_splitter.split_documents(docs)
    print(f"文本分割完毕:文档数据量{len(split_docs)}条")
    # 3 获取嵌入模型
    chat = ChatModel()
    embedding = chat.get_embedding_model()
    #  4 创建neo4j存储
    vector_store = Neo4jVector.from_documents(
        documents=split_docs,
        embedding=embedding,
        url=os.getenv("NEO4J_URL"),
        username=os.getenv("NEO4J_USERNAME"),
        password=os.getenv("NEO4J_PASSWORD"),
        index_name="企业",  # 索引名称
        node_label="企业",  # 节点标签
        text_node_property="text",  # 节点属性
        embedding_node_property="embedding"
    )
    print(" 数据已成功存入Neo4j向量索引!")
